import torch
import os


def checkpoint_remove_opimizer(model_path: str, output_path: str = None):
    """Remove the opimizer parameters in a checkpoint

    Parameters
    ----------
    model_path : str
        The path of checkpoint, e.g., "/NAS/models/sr/your_model.pth"
    output_path : str, optional
        The path of rpocessed checkpoint, e.g., "/NAS/models/sr/your_model2.pth"
        By default the path is the same with the model_path but adding "-para.pth" in the file name.
    """

    if not model_path.endswith(".pth"):
        raise ValueError("The file name should be in *.pth and trained from PyTorch.")

    try:
        file_size = os.path.getsize(model_path)
        file_size_mb = file_size / 1000000
    except FileNotFoundError as e:
        print(f"Read the file with ERROR {e}")

    checkpoint = torch.load(model_path, map_location="cpu")
    # remove optimizer for smaller file size
    if "optimizer" in checkpoint:
        del checkpoint["optimizer"]
    # if it is necessary to remove some sensitive data in checkpoint['meta'],
    # add the code here.

    if output_path is None:
        output_path = model_path.rstrip(".pth") + "-para.pth"

    torch.save(checkpoint, output_path)
    file_size2 = os.path.getsize(output_path)
    file_size2_mb = file_size2 / 1000000

    print(f"Successfully exporting checkpoint {output_path}")
    print(f"The size of '{model_path}' is {file_size_mb:.2f} MB.")
    print(f"The size of '{output_path}' is {file_size2_mb:.2f} MB.")


if __name__ == "__main__":
    checkpoint_remove_opimizer(
        model_path="/NAS/models/cloud_segmentation/SWIN/20230505_swin_s-ssl_4c.pth",
        output_path="/NAS6/Members/larrylai/20230505_swin-s-ssl_4c-para.pth",
    )
